{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import numpy as np\n",
    "import os\n",
    "from tqdm import tqdm\n",
    "import json\n",
    "import shutil\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import xml.etree.ElementTree as ET\n",
    "from xml import etree\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mkdir(path):\n",
    "    if not os.path.exists(path):\n",
    "        os.makedirs(path)\n",
    "        \n",
    "def xml_init(xml_path):\n",
    "    tree = ET.parse(xml_path)\n",
    "    root = tree.getroot() \n",
    "    # 移除所有image_node\n",
    "    for image_node in root.findall(\"image\"):\n",
    "        root.remove(image_node)\n",
    "    return tree,root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_xml_image(xml_root,cvatBasePath,image_file_path,box_list,idx):\n",
    "    height,width,_ = cv2.imread(cvatBasePath+image_file_path).shape\n",
    "    image_node = ET.SubElement(xml_root,\"image\",{'id':str(idx),\"name\":str(image_file_path),\"width\":str(width),\"height\":str(height)})\n",
    "    n = len(box_list)\n",
    "    for i in range(n):\n",
    "        box = box_list[i]\n",
    "        box_node = ET.SubElement(image_node,\"box\",{\"label\":box[0],\"occluded\":\"0\",\"source\":\"manual\",\"xtl\":str(box[1]),\"ytl\":str(box[2]),\"xbr\":str(box[3]),\"ybr\":str(box[4]),\"z_order\":\"0\"})\n",
    "        box_node.tail = \"\\n   \"\n",
    "    image_node.tail = \"\\n  \"\n",
    "    idx += 1\n",
    "    return xml_root,idx\n",
    "\n",
    "def update_lable_json(label_path,image_dir,all_class_dict):\n",
    "    mkdir(label_path)\n",
    "    jsontxt = []\n",
    "    for k in all_class_dict:\n",
    "        for i in range(all_class_dict[k]):\n",
    "            json_dict = {}\n",
    "            json_dict[\"name\"] = k+\"_\"+str(i+1)\n",
    "            json_dict[\"color\"] =\"#fa3253\"\n",
    "            json_dict[\"attributes\"] =[]\n",
    "            jsontxt.append(json_dict)\n",
    "    with open(os.path.join(label_path,\"{}.json\".format(image_dir)),\"w\") as f:\n",
    "        json.dump(jsontxt, f,indent=6)\n",
    "        \n",
    "\n",
    "def update_map_relation(path,name,map_list):\n",
    "    mkdir(path)\n",
    "    with open(os.path.join(path,name+'.txt'), 'a') as f:\n",
    "        for i,cur_map in enumerate(map_list):\n",
    "            if i != 0:\n",
    "                f.write(\"\\n\")\n",
    "            f.write(str(cur_map))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_labels(label_file_path,all_class_dict):\n",
    "\n",
    "    # read txt\n",
    "    with open(label_file_path,\"r\") as f:\n",
    "        labels_txt = f.readlines()\n",
    "    labels = []\n",
    "    for lt in labels_txt:\n",
    "        label_txt = lt[:-1].split(\" \")\n",
    "        label = [label_txt[0]] + [int(label_txt[1]) ]+ [ float(l) for l in label_txt[2:]]\n",
    "        labels.append(label)\n",
    "\n",
    "    # init\n",
    "    cur_class_dict = dict()\n",
    "    box_list = list()\n",
    "    lidar_map_cvat_list = list()\n",
    "    \n",
    "    # start\n",
    "    for obj in labels:\n",
    "        class_name = obj[0]\n",
    "        if class_name not in cur_class_dict:\n",
    "            cur_class_dict[class_name] = 1\n",
    "        else:\n",
    "            cur_class_dict[class_name] += 1\n",
    "        \n",
    "        lidar_obj_name = class_name + \"_\" + str(obj[1])\n",
    "        cvat_obj_name = class_name + \"_\" + str(cur_class_dict[class_name])\n",
    "        lidar_map_cvat_list.append([lidar_obj_name,cvat_obj_name])\n",
    "        \n",
    "        # label    [class_name,obj_id,bbox_xy0[0],bbox_xy0[1],bbox_xy1[0],bbox_xy1[1]]+dim+location+[rotation]\n",
    "        [xbr,ybr,xtl,ytl,]=[round(obj[2]),round(obj[3]),round(obj[4]),round(obj[5])]\n",
    "        if  abs(xbr-xtl)>2000 or (ybr-ytl)>2000:\n",
    "            continue\n",
    "        # xbr = round(obj[2])\n",
    "        # xtl = round(obj[3])\n",
    "        # ybr = round(obj[4])\n",
    "        # ytl = round(obj[5])\n",
    "        \n",
    "        \n",
    "        box_list.append([cvat_obj_name,xtl,ytl,xbr,ybr])\n",
    "        \n",
    "        # 更新all_class_dict的最大值\n",
    "        for cur_class in cur_class_dict:\n",
    "            if cur_class not in all_class_dict:\n",
    "                all_class_dict[cur_class] = cur_class_dict[cur_class]\n",
    "            elif all_class_dict[cur_class] < cur_class_dict[cur_class]:\n",
    "                all_class_dict[cur_class] = cur_class_dict[cur_class]\n",
    "        \n",
    "    return box_list,lidar_map_cvat_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trans_to_cvat(dataPath,cvatBasePath,imageBaseDir,targetPath):\n",
    "    labels_files_path =os.path.join(dataPath,\"labels\")\n",
    "    labels_files =  os.listdir(labels_files_path)\n",
    "    all_class_dict = dict()\n",
    "    image_dirs_list = os.listdir(cvatBasePath+imageBaseDir)\n",
    "    images_map = {}\n",
    "    # global image_dir,file_name,xml_root,xml_tree\n",
    "\n",
    "    for image_dir in image_dirs_list:\n",
    "        # if not image_dir == \"left_back\":\n",
    "        #     continue\n",
    "        \n",
    "        image_dir_path = os.path.join(cvatBasePath+imageBaseDir,image_dir)\n",
    "        if not os.path.isdir(image_dir_path):\n",
    "            continue\n",
    "        images_list = os.listdir(image_dir_path)\n",
    "        names_list = [p.split(\".\")[0] for p in images_list]\n",
    "        xml_tree,xml_root= xml_init(os.path.join(targetPath,\"{}.xml\".format(image_dir)))\n",
    "        print(image_dir)\n",
    "        bar = tqdm(enumerate(names_list))\n",
    "        images_map[image_dir]={}\n",
    "        for idx,file_name in bar:\n",
    "            label_file = file_name + \".txt\"\n",
    "            label_file_path = os.path.join(labels_files_path,label_file)\n",
    "            bar.set_description(\"[{}/{}]\".format(idx,len(names_list)))\n",
    "            # global box_list,lidar_map_cvat_list\n",
    "            # box_list,lidar_map_cvat_list = get_labels(label_file_path,all_class_dict)\n",
    "            image_file_path =os.path.join(imageBaseDir,image_dir,file_name)+\".png\"\n",
    "            images_map[image_dir][idx] = {image_file_path}\n",
    "    #         xml_root,idx = update_xml_image(xml_root,cvatBasePath,image_file_path,box_list,idx)\n",
    "    #         update_map_relation(os.path.join(targetPath,\"map_relation\",image_dir),file_name,lidar_map_cvat_list)\n",
    "    #     xml_tree.write(os.path.join(targetPath,\"{}_convert.xml\".format(image_dir)),encoding=\"utf-8\",xml_declaration=True)\n",
    "    # update_lable_json(targetPath,\"label\",all_class_dict)\n",
    "    np.save(os.path.join(targetPath,\"images_map.dict\"),images_map) \n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "right_back\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[473/474]: : 474it [00:00, 1762.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "front_middle\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[1306/1307]: : 1307it [00:00, 1804.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "left_back\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[356/357]: : 357it [00:00, 1795.05it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "right_front\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[761/762]: : 762it [00:00, 1887.14it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "left_front\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[810/811]: : 811it [00:00, 1763.84it/s]\n"
     ]
    }
   ],
   "source": [
    "labelPath = \"/home/fuyu/code/convert_tools/boleidun\"\n",
    "cvatBasePath = \"/home/fuyu/data/cvat/\"\n",
    "imageBaseDir = \"boleidun25d\"\n",
    "targetPath = \"/home/fuyu/code/convert_tools/boleidun_cvat\"\n",
    "trans_to_cvat(labelPath,cvatBasePath,imageBaseDir,targetPath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "p = os.path.join(targetPath,\"images_map.dict.npy\")\n",
    "d = np.load(p,allow_pickle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'boleidun25d/left_back/2022-09-05-16-07-58_3_left_back_0052.png'}"
      ]
     },
     "execution_count": 60,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d.item()[\"left_back\"][263]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'boleidun25d/left_front/2022-09-05-15-01-17_2_left_back_0010.png'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "image_file_path =os.path.join(imageBaseDir,image_dir,file_name)+\".png\"\n",
    "image_file_path\n",
    "# \"boleidun25d/front_middle/2022-09-05-15-00-54_0_front_middle_0005.png\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# idx = 0\n",
    "# xml_root,idx = update_xml_image(xml_root,image_file_path,box_list,idx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "e1a1b001e220d928b91f8c0266777d646605be3640513aa809eb904b5d0cdf9a"
  },
  "kernelspec": {
   "display_name": "Python 3.8.13 ('tools')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
